import pymysql
import json

def test_connection(host, port, username, password, database):
    """测试数据库连接"""
    try:
        conn = pymysql.connect(
            host=host,
            port=int(port),
            user=username,
            password=password,
            database=database,
            connect_timeout=5
        )
        conn.close()
        return True, "连接成功"
    except Exception as e:
        return False, str(e)

def get_tables(host, port, username, password, database):
    """获取数据库中的所有表"""
    try:
        conn = pymysql.connect(
            host=host,
            port=int(port),
            user=username,
            password=password,
            database=database
        )
        cursor = conn.cursor()
        cursor.execute("SHOW TABLES")
        tables = [table[0] for table in cursor.fetchall()]
        cursor.close()
        conn.close()
        return tables
    except Exception as e:
        return []

def get_table_schema(host, port, username, password, database, table_name):
    """获取表结构"""
    try:
        conn = pymysql.connect(
            host=host,
            port=int(port),
            user=username,
            password=password,
            database=database
        )
        cursor = conn.cursor(pymysql.cursors.DictCursor)
        
        # 获取表结构
        cursor.execute(f"DESCRIBE `{table_name}`")
        columns = cursor.fetchall()
        
        # 获取主键
        cursor.execute(f"""
            SELECT COLUMN_NAME
            FROM INFORMATION_SCHEMA.KEY_COLUMN_USAGE
            WHERE TABLE_SCHEMA = %s
            AND TABLE_NAME = %s
            AND CONSTRAINT_NAME = 'PRIMARY'
        """, (database, table_name))
        primary_keys = [row['COLUMN_NAME'] for row in cursor.fetchall()]
        
        # 处理结果
        schema = []
        for column in columns:
            schema.append({
                'name': column['Field'],
                'type': column['Type'],
                'nullable': column['Null'] == 'YES',
                'default': column['Default'],
                'is_primary': column['Field'] in primary_keys
            })
        
        cursor.close()
        conn.close()
        return schema
    except Exception as e:
        return []

def get_table_data_sample(host, port, username, password, database, table_name, limit=10):
    """获取表数据样本"""
    try:
        conn = pymysql.connect(
            host=host,
            port=int(port),
            user=username,
            password=password,
            database=database
        )
        cursor = conn.cursor(pymysql.cursors.DictCursor)
        
        # 获取表数据
        cursor.execute(f"SELECT * FROM `{table_name}` LIMIT {limit}")
        data = cursor.fetchall()
        
        # 处理结果
        result = []
        for row in data:
            processed_row = {}
            for key, value in row.items():
                # 处理不可序列化的数据类型
                if isinstance(value, bytes):
                    processed_row[key] = value.hex()
                else:
                    processed_row[key] = value
            result.append(processed_row)
        
        cursor.close()
        conn.close()
        return result
    except Exception as e:
        return []